#!/bin/bash
# -*- coding: utf-8 -*-
# BentoML Multi-Model Service Startup Script
# Compatible with multiple Linux distributions
# Supports: Ubuntu, Debian, CentOS, RHEL, Fedora, Arch, etc.

set -e  # Exit on error

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
CONDA_ENV_NAME="bentoml"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MINICONDA_DIR="$HOME/miniconda3"
PYTHON_VERSION="3.10"
SERVICE_PORT="${BENTOML_PORT:-3000}"
SERVICE_HOST="${BENTOML_HOST:-0.0.0.0}"

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Print banner
print_banner() {
    echo -e "${GREEN}"
    echo "=========================================="
    echo "  BentoML Multi-Model Service Launcher"
    echo "=========================================="
    echo -e "${NC}"
    echo "Models:"
    echo "  - m3e-large (Embedding)"
    echo "  - bge-reranker-v2-m3 (Reranker)"
    echo "  - PaddleOCR v4 (OCR)"
    echo ""
}

# Detect Linux distribution
detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        DISTRO=$ID
        DISTRO_VERSION=$VERSION_ID
    elif [ -f /etc/lsb-release ]; then
        . /etc/lsb-release
        DISTRO=$DISTRIB_ID
        DISTRO_VERSION=$DISTRIB_RELEASE
    else
        DISTRO="unknown"
        DISTRO_VERSION="unknown"
    fi
    
    log_info "Detected distribution: $DISTRO $DISTRO_VERSION"
}

# Check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Install system dependencies based on distribution
install_system_dependencies() {
    log_info "Checking system dependencies..."
    
    local packages_to_install=()
    
    # Check for required packages
    if ! command_exists wget && ! command_exists curl; then
        packages_to_install+=("wget")
    fi
    
    if ! command_exists git; then
        packages_to_install+=("git")
    fi
    
    # Check for build essentials
    if ! command_exists gcc; then
        case $DISTRO in
            ubuntu|debian)
                packages_to_install+=("build-essential")
                ;;
            centos|rhel|fedora)
                packages_to_install+=("gcc" "gcc-c++" "make")
                ;;
            arch|manjaro)
                packages_to_install+=("base-devel")
                ;;
        esac
    fi
    
    # Install packages if needed
    if [ ${#packages_to_install[@]} -gt 0 ]; then
        log_info "Installing system dependencies: ${packages_to_install[*]}"
        
        case $DISTRO in
            ubuntu|debian)
                sudo apt-get update
                sudo apt-get install -y "${packages_to_install[@]}"
                # Additional libraries for PaddleOCR
                sudo apt-get install -y libgomp1 libglib2.0-0 libsm6 libxext6 libxrender-dev libgl1-mesa-glx
                ;;
            centos|rhel)
                sudo yum install -y "${packages_to_install[@]}"
                sudo yum install -y libgomp glib2 libSM libXext libXrender mesa-libGL
                ;;
            fedora)
                sudo dnf install -y "${packages_to_install[@]}"
                sudo dnf install -y libgomp glib2 libSM libXext libXrender mesa-libGL
                ;;
            arch|manjaro)
                sudo pacman -Sy --noconfirm "${packages_to_install[@]}"
                sudo pacman -Sy --noconfirm libgomp glib2 libsm libxext libxrender mesa
                ;;
            *)
                log_warning "Unknown distribution. Please install dependencies manually: ${packages_to_install[*]}"
                ;;
        esac
        
        log_success "System dependencies installed"
    else
        log_success "All system dependencies are already installed"
    fi
}

# Check and install Miniconda
check_install_conda() {
    log_info "Checking for Conda installation..."
    
    if command_exists conda; then
        log_success "Conda is already installed"
        # Initialize conda for bash if not already done
        if [ -f "$HOME/miniconda3/etc/profile.d/conda.sh" ]; then
            source "$HOME/miniconda3/etc/profile.d/conda.sh"
        elif [ -f "$HOME/anaconda3/etc/profile.d/conda.sh" ]; then
            source "$HOME/anaconda3/etc/profile.d/conda.sh"
        fi
        return 0
    fi
    
    log_info "Conda not found. Installing Miniconda..."
    
    # Detect architecture
    ARCH=$(uname -m)
    case $ARCH in
        x86_64)
            MINICONDA_ARCH="x86_64"
            ;;
        aarch64|arm64)
            MINICONDA_ARCH="aarch64"
            ;;
        *)
            log_error "Unsupported architecture: $ARCH"
            exit 1
            ;;
    esac
    
    # Download Miniconda
    MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh"
    MINICONDA_INSTALLER="/tmp/miniconda_installer.sh"
    
    log_info "Downloading Miniconda from $MINICONDA_URL"
    
    if command_exists wget; then
        wget -O "$MINICONDA_INSTALLER" "$MINICONDA_URL"
    elif command_exists curl; then
        curl -L -o "$MINICONDA_INSTALLER" "$MINICONDA_URL"
    else
        log_error "Neither wget nor curl is available. Cannot download Miniconda."
        exit 1
    fi
    
    # Install Miniconda
    log_info "Installing Miniconda to $MINICONDA_DIR"
    bash "$MINICONDA_INSTALLER" -b -p "$MINICONDA_DIR"
    
    # Initialize conda
    "$MINICONDA_DIR/bin/conda" init bash
    source "$MINICONDA_DIR/etc/profile.d/conda.sh"
    
    # Clean up
    rm "$MINICONDA_INSTALLER"
    
    log_success "Miniconda installed successfully"
}

# Create or activate conda environment
setup_conda_environment() {
    log_info "Setting up Conda environment: $CONDA_ENV_NAME"
    
    # Check if environment exists
    if conda env list | grep -q "^$CONDA_ENV_NAME "; then
        log_success "Environment '$CONDA_ENV_NAME' already exists, activating it..."
        conda activate "$CONDA_ENV_NAME"
        
        # Check if this is a valid Python environment
        if python --version >/dev/null 2>&1; then
            CURRENT_PYTHON_VERSION=$(python --version 2>&1 | cut -d' ' -f2 | cut -d'.' -f1,2)
            log_info "Current Python version in environment: $CURRENT_PYTHON_VERSION"
            
            # Check if Python version matches our requirement
            if [[ "$CURRENT_PYTHON_VERSION" == "$PYTHON_VERSION" ]]; then
                log_success "Python version matches requirement ($PYTHON_VERSION)"
                return 0
            else
                log_warning "Python version mismatch. Required: $PYTHON_VERSION, Found: $CURRENT_PYTHON_VERSION"
                read -p "Do you want to recreate the environment with correct Python version? (y/N): " -n 1 -r
                echo
                if [[ $REPLY =~ ^[Yy]$ ]]; then
                    log_info "Removing existing environment..."
                    conda deactivate 2>/dev/null || true
                    conda env remove -n "$CONDA_ENV_NAME" -y
                    log_info "Creating new environment with Python $PYTHON_VERSION..."
                    conda create -n "$CONDA_ENV_NAME" python=$PYTHON_VERSION -y
                else
                    log_info "Continuing with existing Python version..."
                    return 0
                fi
            fi
        else
            log_warning "Environment exists but Python is not accessible"
            read -p "Do you want to recreate the environment? (y/N): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                log_info "Removing existing environment..."
                conda deactivate 2>/dev/null || true
                conda env remove -n "$CONDA_ENV_NAME" -y
                log_info "Creating new environment..."
                conda create -n "$CONDA_ENV_NAME" python=$PYTHON_VERSION -y
            else
                log_error "Cannot proceed with broken environment"
                exit 1
            fi
        fi
    else
        log_info "Environment '$CONDA_ENV_NAME' does not exist, creating new one..."
        conda create -n "$CONDA_ENV_NAME" python=$PYTHON_VERSION -y
    fi
    
    log_success "Conda environment ready"
}

# Install Python dependencies
install_python_dependencies() {
    log_info "Checking Python dependencies..."
    
    # Ensure conda environment is activated
    conda activate "$CONDA_ENV_NAME"
    
    # Check if key dependencies are already installed
    log_info "Checking for existing dependencies..."
    
    MISSING_DEPS=()
    KEY_PACKAGES=("bentoml" "torch" "transformers" "paddleocr" "sentence-transformers" "FlagEmbedding")
    
    for package in "${KEY_PACKAGES[@]}"; do
        if ! python -c "import $package" >/dev/null 2>&1; then
            MISSING_DEPS+=("$package")
        fi
    done
    
    if [ ${#MISSING_DEPS[@]} -eq 0 ]; then
        log_success "All key dependencies are already installed"
        
        # Still verify critical packages
        log_info "Verifying critical package versions..."
        python -c "
import bentoml; print('BentoML version:', bentoml.__version__)
import torch; print('PyTorch version:', torch.__version__)
try:
    import modelscope; print('ModelScope SDK version:', modelscope.__version__)
except ImportError:
    print('ModelScope SDK: Not installed')
" || {
            log_warning "Some packages may need updates"
        }
        
        # Ask if user wants to update dependencies
        read -p "Dependencies are installed. Do you want to update them? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Skipping dependency installation"
            return 0
        fi
    else
        log_info "Missing dependencies detected: ${MISSING_DEPS[*]}"
    fi
    
    # Upgrade pip
    log_info "Upgrading pip..."
    pip install --upgrade pip
    
    # Install dependencies from requirements file
    if [ -f "$SCRIPT_DIR/requirements_bentoml.txt" ]; then
        log_info "Installing/updating from requirements_bentoml.txt..."
        pip install -r "$SCRIPT_DIR/requirements_bentoml.txt" -i https://pypi.tuna.tsinghua.edu.cn/simple
        
        # Verify ModelScope SDK installation
        log_info "Verifying ModelScope SDK installation..."
        python -c "import modelscope; print('ModelScope SDK version:', modelscope.__version__)" || {
            log_warning "ModelScope SDK verification failed, but continuing..."
        }
    else
        log_error "requirements_bentoml.txt not found in $SCRIPT_DIR"
        exit 1
    fi
    
    log_success "Python dependencies installed/updated"
}

# Download models
download_models() {
    log_info "Checking model availability..."
    
    # Activate conda environment
    conda activate "$CONDA_ENV_NAME"
    
    # Check if models are already cached in default directories
    log_info "Checking for existing models in cache..."
    
    MODELS_AVAILABLE=true
    
    # Quick check for model availability in default cache directories
    python -c "
import os
import sys
from pathlib import Path

print('Checking models in default cache directories...')

# Try to import and test model loading
try:
    from sentence_transformers import SentenceTransformer
    from FlagEmbedding import FlagReranker
    
    # Quick test - this will check if models are cached
    print('Testing model imports...')
    
    # This is just a quick import test, not actual loading
    print('✓ Model libraries available')
    print('Models will be downloaded automatically if not cached')
    sys.exit(0)
    
except ImportError as e:
    print(f'Model libraries not available: {e}')
    sys.exit(1)
except Exception as e:
    print(f'Error checking models: {e}')
    sys.exit(1)
" && {
        log_success "Models appear to be already available"
        
        # Ask if user wants to re-download models
        read -p "Models seem to be cached. Do you want to re-download/verify them? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Skipping model download"
            return 0
        fi
    } || {
        log_info "Models need to be downloaded"
    }
    
    # Run model download script
    if [ -f "$SCRIPT_DIR/download_models.py" ]; then
        log_info "Running model download script..."
        python "$SCRIPT_DIR/download_models.py"
        
        if [ $? -eq 0 ]; then
            log_success "Models downloaded successfully"
        else
            log_warning "Some models may have failed to download. Check logs above."
            read -p "Continue anyway? (y/N): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                exit 1
            fi
        fi
    else
        log_error "download_models.py not found in $SCRIPT_DIR"
        exit 1
    fi
}

# Start BentoML service
start_service() {
    log_info "Starting BentoML service..."
    
    # Activate conda environment
    conda activate "$CONDA_ENV_NAME"
    
    # Change to script directory
    cd "$SCRIPT_DIR"
    
    # Check if service file exists
    if [ ! -f "$SCRIPT_DIR/bentoml_service.py" ]; then
        log_error "bentoml_service.py not found in $SCRIPT_DIR"
        exit 1
    fi
    
    log_info "Service configuration:"
    log_info "  - Host: $SERVICE_HOST"
    log_info "  - Port: $SERVICE_PORT"
    log_info "  - Working directory: $SCRIPT_DIR"
    
    # Start the service
    log_success "Starting BentoML service..."
    echo ""
    echo -e "${GREEN}=========================================="
    echo "  Service is starting..."
    echo "  Access endpoints at:"
    echo "  - http://${SERVICE_HOST}:${SERVICE_PORT}"
    echo "==========================================${NC}"
    echo ""
    
    # Run BentoML serve
    bentoml serve bentoml_service:MultiModelService \
        --host "$SERVICE_HOST" \
        --port "$SERVICE_PORT" \
        --reload
}

# Main execution flow
main() {
    print_banner
    
    # Detect distribution
    detect_distro
    
    # Check if running as root (not recommended)
    if [ "$EUID" -eq 0 ]; then
        log_warning "Running as root is not recommended"
        read -p "Continue anyway? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    fi
    
    # Step 1: Install system dependencies
    log_info "Step 1/6: Installing system dependencies..."
    install_system_dependencies
    echo ""
    
    # Step 2: Check and install Conda
    log_info "Step 2/6: Checking Conda installation..."
    check_install_conda
    echo ""
    
    # Step 3: Setup Conda environment
    log_info "Step 3/6: Setting up Conda environment..."
    setup_conda_environment
    echo ""
    
    # Step 4: Install Python dependencies
    log_info "Step 4/6: Installing Python dependencies..."
    install_python_dependencies
    echo ""
    
    # Step 5: Download models
    log_info "Step 5/6: Downloading models..."
    download_models
    echo ""
    
    # Step 6: Start service
    log_info "Step 6/6: Starting BentoML service..."
    start_service
}

# Run main function
main
